package com.study.recommendation

import org.apache.spark.ml.evaluation.RegressionEvaluator
import org.apache.spark.ml.recommendation.ALS
import org.apache.spark.sql.SparkSession

/**
 * 协同过滤
 * 使用ALS(alternating least squares/交替最小二乘)
 *
 * @author stephen
 * @date 2019-08-29 18:21
 */
case class Rating(userId: Int, movieId: Int, rating: Float)

object CollaborativeFilteringDemo {

  def main(args: Array[String]): Unit = {

    val spark = SparkSession.builder()
      .appName(this.getClass.getSimpleName)
      .master("local[*]")
      .getOrCreate()

    spark.sparkContext.setLogLevel("warn")

    import spark.implicits._

    val path = this.getClass.getClassLoader.getResource("data/mllib/sample_movielens_data.txt").getPath
    val ratings = spark.read.textFile(path)
      .map[Rating]((row: String) => {
        val fields = row.split("::")
        assert(fields.size > 2)
        Rating(fields(0).toInt, fields(1).toInt, fields(2).toFloat)
      })
      .toDF

    // 拆分成训练集和测试集
    val Array(training, test) = ratings.randomSplit(Array(0.8, 0.2))

    // 建立ALS交替最小二乘算法模型并训练
    val als = new ALS()
      .setMaxIter(5)
      .setRegParam(0.01)
      .setUserCol("userId")
      .setItemCol("movieId")
      .setRatingCol("rating")
    val model = als.fit(training)

    // 使用drop的冷启动策略后，未在训练集中出现的都不会给出预测结果。
    model.setColdStartStrategy("drop")
    val predictions = model.transform(test)
    predictions.show(false)

    val evaluator = new RegressionEvaluator()
      .setMetricName("rmse")
      .setLabelCol("rating")
      .setPredictionCol("prediction")
    val rmse = evaluator.evaluate(predictions)
    println(s"Root-mean-square error = $rmse")

    // Generate top 10 movie recommendations for each user
    val userRecs = model.recommendForAllUsers(10)
    //userRecs.show(false)
    // Generate top 10 user recommendations for each movie
    val movieRecs = model.recommendForAllItems(10)

    // Generate top 10 movie recommendations for a specified set of users
    val users = ratings.select(als.getUserCol).distinct().limit(3)
    val userSubsetRecs = model.recommendForUserSubset(users, 10)
    // Generate top 10 user recommendations for a specified set of movies
    val movies = ratings.select(als.getItemCol).distinct().limit(3)
    val movieSubSetRecs = model.recommendForItemSubset(movies, 10)


  }
}
